import tensorflow as tf
from tensorflow.keras import layers, models
from tensorflow.keras.datasets import cifar10
from tensorflow.keras.utils import to_categorical
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.callbacks import EarlyStopping, ReduceLROnPlateau
import numpy as np
from sklearn.metrics import classification_report, confusion_matrix

# 超参数配置
BATCH_SIZE = 128  # 批量大小
EPOCHS = 50       # 训练轮数
LEARNING_RATE = 1e-3  # 学习率
IMG_SIZE = 32     # 图像大小

# 数据加载与预处理
def load_data():
    # 加载CIFAR-10数据集
    (train_images, train_labels), (test_images, test_labels) = cifar10.load_data()
    
    # 数据预处理：归一化到[0, 1]范围
    train_images = train_images.astype('float32') / 255.0
    test_images = test_images.astype('float32') / 255.0
    
    # 将标签进行One-hot编码
    train_labels = to_categorical(train_labels, 10)
    test_labels = to_categorical(test_labels, 10)
    
    return (train_images, train_labels), (test_images, test_labels)

# 添加残差模块
def residual_block(x, filters):
    shortcut = x
    if shortcut.shape[-1] != filters:
        shortcut = layers.Conv2D(filters, (1, 1), strides=1, padding='same')(shortcut)
        shortcut = layers.BatchNormalization()(shortcut)
    
    x = layers.Conv2D(filters, (3, 3), activation='relu', padding='same')(x)
    x = layers.BatchNormalization()(x)
    x = layers.Conv2D(filters, (3, 3), activation='relu', padding='same')(x)
    x = layers.BatchNormalization()(x)
    x = layers.add([shortcut, x])
    x = layers.Activation('relu')(x)
    return x

# 改进的CNN模型
def build_model():
    inputs = layers.Input(shape=(IMG_SIZE, IMG_SIZE, 3))
    
    # 数据增强层（仅在训练时激活）
    x = layers.RandomRotation(0.1)(inputs)  # 随机旋转
    x = layers.RandomZoom(0.1)(x)           # 随机缩放
    x = layers.RandomFlip("horizontal")(x)   # 随机水平翻转
    
    # 初始卷积层
    x = layers.Conv2D(64, (3, 3), activation='relu', padding='same')(x)
    x = layers.BatchNormalization()(x)
    
    # 第一个残差模块
    x = residual_block(x, 64)
    x = layers.MaxPooling2D((2, 2))(x)
    x = layers.Dropout(0.25)(x)
    
    # 第二个残差模块
    x = residual_block(x, 128)
    x = layers.MaxPooling2D((2, 2))(x)
    x = layers.Dropout(0.25)(x)
    
    # 第三个残差模块
    x = residual_block(x, 256)
    x = layers.MaxPooling2D((2, 2))(x)
    x = layers.Dropout(0.25)(x)
    
    # 全连接层
    x = layers.Flatten()(x)
    x = layers.Dense(512, activation='relu')(x)
    x = layers.BatchNormalization()(x)
    x = layers.Dropout(0.5)(x)
    outputs = layers.Dense(10, activation='softmax')(x)
    
    model = models.Model(inputs, outputs)
    
    # 定义优化器
    optimizer = tf.keras.optimizers.Adam(learning_rate=LEARNING_RATE)
    # 编译模型
    model.compile(optimizer=optimizer,
                  loss='categorical_crossentropy',
                  metrics=['accuracy'])
    return model

# 训练流程
def main():
    # 加载数据
    (train_images, train_labels), (test_images, test_labels) = load_data()
    
    # 创建模型
    model = build_model()
    model.summary()  # 打印模型结构
    
    # 定义回调函数
    early_stop = EarlyStopping(monitor='val_loss', patience=15, restore_best_weights=True)  # 早停
    reduce_lr = ReduceLROnPlateau(monitor='val_loss', factor=0.2, patience=5, min_lr=1e-6)  # 学习率衰减
    
    # 数据增强生成器（生成新样本，增强泛化能力）
    train_datagen = ImageDataGenerator(
        rotation_range=15,  # 随机旋转角度范围
        width_shift_range=0.1,  # 水平平移范围
        height_shift_range=0.1,  # 垂直平移范围
        horizontal_flip=True,  # 水平翻转
        zoom_range=0.2  # 缩放范围
    )
    
    # 训练模型
    history = model.fit(
        train_datagen.flow(train_images, train_labels, batch_size=BATCH_SIZE),  # 数据增强后的训练数据
        epochs=EPOCHS,
        validation_data=(test_images, test_labels),  # 验证数据
        callbacks=[early_stop, reduce_lr],  # 回调函数
        verbose=1  # 输出训练进度
    )
    
    # 保存模型
    model.save('improved_classifier_model.h5')
    
    # 评估模型
    print("\n评估结果:")
    test_loss, test_acc = model.evaluate(test_images, test_labels)
    print(f"测试准确率: {test_acc:.4f}")
    
    # 生成分类报告
    y_pred = model.predict(test_images)
    y_pred_classes = np.argmax(y_pred, axis=1)
    y_true = np.argmax(test_labels, axis=1)
    
    print("\n分类报告:")
    print(classification_report(y_true, y_pred_classes))

if __name__ == "__main__":
    main()